#include <gtest/gtest.h>

#include "trt_define.h"
#include "utils/logger.h"
#include "utils/file_utils.h"
#include "utils/nv_utils.h"
#include "trt_engine/trt_engine.h"
#include "trt_engine/preprocess_kernel.cuh"
#include "trt_engine/monopoly_allocator.h"

_TRT_INFER_BEGIN

class EngineTest : public ::testing::Test
{
public:
    std::string source_model_path;
    std::string save_model_path;
    std::string image_path;

protected:
    EngineTest() {
        LOGER_INST.init();
        LOGI << "test_tensor";
        source_model_path = "D:\\learnSpace\\trt_infer\\test_data\\classifier.onnx";
        save_model_path = "D:\\learnSpace\\trt_infer\\test_data\\classifier.engine";
        image_path = "D:\\learnSpace\\trt_infer\\test_data\\dog.jpg";
    }
    ~EngineTest() override {}
    void SetUp() override {}
    void TearDown() override {}
};

TEST_F(EngineTest, test_tensor)
{
    auto engine = load_infer(save_model_path);
    ASSERT_TRUE(engine) << "load engine failed";

    engine->print();

    auto input = engine->input();
    auto output = engine->output();
    int input_width = input->width();
    int input_height = input->height();

    auto image = cv::imread(image_path);
    ASSERT_NE(image.data, nullptr) << "read image failed";

    cv::resize(image, image, cv::Size(input_width, input_height));
    image.convertTo(image, CV_32F);

    cv::Mat channel_based[3];
    for (int i = 0; i < 3; ++i)
        channel_based[i] = cv::Mat(input_height, input_width, CV_32F, input->cpu<float>(0, 2 - i));

    float mean[] = {0.406, 0.456, 0.485};
    float std[] = {0.225, 0.224, 0.229};
    cv::split(image, channel_based);
    for (int i = 0; i < 3; ++i)
        channel_based[i] = (channel_based[i] / 255.0f - mean[i]) / std[i];

    engine->forward(true);
    int num_classes   = output->size(1);
    float* prob       = output->cpu<float>();
    int predict_label = std::max_element(prob, prob + num_classes) - prob;
    float confidence  = prob[predict_label];
    LOGI << "predict label: " << predict_label << ", confidence: " << confidence;
}

TEST_F(EngineTest, MonopolyAllocator)
{
    MonopolyAllocator<std::string> allocator(2);
    EXPECT_EQ(allocator.num_available(), 2);
    auto data = allocator.query();
    data->data() = std::make_shared<std::string>("hello world");
    EXPECT_EQ(allocator.num_available(), 1);

    auto data2 = allocator.query();
    data2->data() = std::make_shared<std::string>("hello world2");

    data2->release();
    auto data3 = allocator.query();
    EXPECT_NE(data3, nullptr);

    auto data4 = allocator.query();
    EXPECT_EQ(data4, nullptr);
}




_TRT_INFER_END
